#include "Util.h"
#include <windef.h>
#include <ntstrsafe.h>

NTSTATUS Ps2GetProcessFullDosNameByPid(IN HANDLE Pid, IN PUNICODE_STRING lpDosFullName)
{
    HANDLE hFile = NULL;
    ULONG nNeedSize = 0;
    NTSTATUS ntStatus = STATUS_SUCCESS;
    NTSTATUS ntDeviceStatus = STATUS_DEVICE_DOES_NOT_EXIST;
    PEPROCESS lpEprocess = NULL;
    KAPC_STATE ApcState = { 0 };
    PVOID lpBuff = NULL;
    OBJECT_ATTRIBUTES ObjectAttributes = { 0 };
    IO_STATUS_BLOCK IoStatus = { 0 };
    PFILE_OBJECT lpFileObject = NULL;
    PFILE_NAME_INFORMATION lpFileNameInfo = NULL;
    WCHAR wstrFileBuff[MAX_PATH] = { 0 };
    DECLARE_UNICODE_STRING_SIZE(usProcessPath, MAX_PATH);
    DECLARE_UNICODE_STRING_SIZE(usDosDeviceName, MAX_PATH);
    KIRQL Irql;

    PAGED_CODE();

    ntStatus = PsLookupProcessByProcessId(Pid, &lpEprocess);
    if (!NT_SUCCESS(ntStatus))
    {
        return ntStatus;
    }

    do
    {
        KeStackAttachProcess(lpEprocess, &ApcState);

        ntStatus = _ZwQueryInformationProcess(NtCurrentProcess(), ProcessImageFileName, NULL, 0, &nNeedSize);
        if (STATUS_INFO_LENGTH_MISMATCH != ntStatus)
        {
            ERRPRINT(ZwQueryInformationProcess, ntStatus);
            break;
        }

        lpBuff = ExAllocatePool(NonPagedPool, nNeedSize);
        if (!lpBuff)
        {
            ntStatus = STATUS_INSUFFICIENT_RESOURCES;
            ERRPRINT(ExAllocatePool, ntStatus);
            break;
        }

        ntStatus = _ZwQueryInformationProcess(NtCurrentProcess(), ProcessImageFileName, lpBuff, nNeedSize, &nNeedSize);
        if (!NT_SUCCESS(ntStatus))
        {
            ERRPRINT(ZwQueryInformationProcess, ntStatus);
            break;
        }

        Irql = KeGetCurrentIrql();
        if (Irql != PASSIVE_LEVEL)
        {
            ntStatus = STATUS_UNSUCCESSFUL;
            break;
        }

        RtlCopyUnicodeString(&usProcessPath, (PUNICODE_STRING)lpBuff);
        InitializeObjectAttributes(&ObjectAttributes, &usProcessPath, OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE, NULL, NULL);

        ntStatus = ZwCreateFile(&hFile, GENERIC_READ, &ObjectAttributes, &IoStatus, NULL, FILE_ATTRIBUTE_NORMAL, 
            FILE_SHARE_READ | FILE_SHARE_DELETE, FILE_OPEN_IF, FILE_NON_DIRECTORY_FILE | FILE_RANDOM_ACCESS | FILE_SYNCHRONOUS_IO_NONALERT, NULL, 0);
        if (!NT_SUCCESS(ntStatus))
        {
            ERRPRINT(ZwCreateFile, ntStatus);
            break;
        }

        ntStatus = ObReferenceObjectByHandle(hFile, 0, *IoFileObjectType, KernelMode, (PVOID*)&lpFileObject, NULL);
        if (!NT_SUCCESS(ntStatus))
        {
            ERRPRINT(ObReferenceObjectByHandle, ntStatus);
            break;
        }

        lpFileNameInfo = (PFILE_NAME_INFORMATION)wstrFileBuff;
        ntStatus = ZwQueryInformationFile(hFile, &IoStatus, lpFileNameInfo, sizeof(WCHAR) * MAX_PATH, FileNameInformation);
        if (!NT_SUCCESS(ntStatus))
        {
            ERRPRINT(ZwQueryInformationFile, ntStatus);
            break;
        }

        if (!lpFileObject->DeviceObject)
        {
            ntDeviceStatus = STATUS_DEVICE_DOES_NOT_EXIST;
            break;
        }

        ntDeviceStatus = RtlVolumeDeviceToDosName(lpFileObject->DeviceObject, &usDosDeviceName);

    } while (FALSE);

    if (lpFileObject)
        ObDereferenceObject(lpFileObject);

    if (hFile)
        ZwClose(hFile);

    if (lpBuff)
        ExFreePool(lpBuff);

    KeUnstackDetachProcess(&ApcState);

    if (NT_SUCCESS(ntStatus))
    {
        RtlInitUnicodeString(&usProcessPath, lpFileNameInfo->FileName);

        if (NT_SUCCESS(ntDeviceStatus))
        {
            RtlCopyUnicodeString(lpDosFullName, &usDosDeviceName);
            RtlUnicodeStringCat(lpDosFullName, &usProcessPath);
        }
        else
        {
            RtlCopyUnicodeString(lpDosFullName, &usProcessPath);
        }
    }

    if (lpEprocess)
        ObDereferenceObject(lpEprocess);

    return ntStatus;
}

NTSTATUS Ps2GetProcessFullDosName(IN PEPROCESS lpEprocess, IN PUNICODE_STRING lpDosFullName)
{
    HANDLE hFile = NULL;
    ULONG nNeedSize = 0;
    NTSTATUS ntStatus = STATUS_SUCCESS;
    NTSTATUS ntDeviceStatus = STATUS_DEVICE_DOES_NOT_EXIST;
    KAPC_STATE ApcState = { 0 };
    PVOID lpBuff = NULL;
    OBJECT_ATTRIBUTES ObjectAttributes = { 0 };
    IO_STATUS_BLOCK IoStatus = { 0 };
    PFILE_OBJECT lpFileObject = NULL;
    PFILE_NAME_INFORMATION lpFileNameInfo = NULL;
    WCHAR wstrFileBuff[MAX_PATH] = { 0 };
    DECLARE_UNICODE_STRING_SIZE(usProcessPath, MAX_PATH);
    DECLARE_UNICODE_STRING_SIZE(usDosDeviceName, MAX_PATH);
    KIRQL Irql;

    PAGED_CODE();

    if (!lpEprocess)
        return STATUS_UNSUCCESSFUL;

    do
    {
        KeStackAttachProcess(lpEprocess, &ApcState);

        ntStatus = _ZwQueryInformationProcess(NtCurrentProcess(), ProcessImageFileName, NULL, 0, &nNeedSize);
        if (STATUS_INFO_LENGTH_MISMATCH != ntStatus)
        {
            ERRPRINT(ZwQueryInformationProcess, ntStatus);
            break;
        }

        lpBuff = ExAllocatePool(NonPagedPool, nNeedSize);
        if (!lpBuff)
        {
            ntStatus = STATUS_INSUFFICIENT_RESOURCES;
            ERRPRINT(ExAllocatePool, ntStatus);
            break;
        }

        ntStatus = _ZwQueryInformationProcess(NtCurrentProcess(), ProcessImageFileName, lpBuff, nNeedSize, &nNeedSize);
        if (!NT_SUCCESS(ntStatus))
        {
            ERRPRINT(ZwQueryInformationProcess, ntStatus);
            break;
        }

        Irql = KeGetCurrentIrql();
        if (Irql != PASSIVE_LEVEL)
        {
            ntStatus = STATUS_UNSUCCESSFUL;
            break;
        }

        RtlCopyUnicodeString(&usProcessPath, (PUNICODE_STRING)lpBuff);
        InitializeObjectAttributes(&ObjectAttributes, &usProcessPath, OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE, NULL, NULL);

        ntStatus = ZwCreateFile(&hFile, GENERIC_READ, &ObjectAttributes, &IoStatus, NULL, FILE_ATTRIBUTE_NORMAL, 
            FILE_SHARE_READ | FILE_SHARE_DELETE, FILE_OPEN_IF, FILE_NON_DIRECTORY_FILE | FILE_RANDOM_ACCESS | FILE_SYNCHRONOUS_IO_NONALERT, NULL, 0);
        if (!NT_SUCCESS(ntStatus))
        {
            ERRPRINT(ZwCreateFile, ntStatus);
            break;
        }

        ntStatus = ObReferenceObjectByHandle(hFile, 0, *IoFileObjectType, KernelMode, (PVOID*)&lpFileObject, NULL);
        if (!NT_SUCCESS(ntStatus))
        {
            ERRPRINT(ObReferenceObjectByHandle, ntStatus);
            break;
        }

        lpFileNameInfo = (PFILE_NAME_INFORMATION)wstrFileBuff;
        ntStatus = ZwQueryInformationFile(hFile, &IoStatus, lpFileNameInfo, sizeof(WCHAR) * MAX_PATH, FileNameInformation);
        if (!NT_SUCCESS(ntStatus))
        {
            ERRPRINT(ZwQueryInformationFile, ntStatus);
            break;
        }

        if (!lpFileObject->DeviceObject)
        {
            ntDeviceStatus = STATUS_DEVICE_DOES_NOT_EXIST;
            break;
        }

        ntDeviceStatus = RtlVolumeDeviceToDosName(lpFileObject->DeviceObject, &usDosDeviceName);

    } while (FALSE);

    if (lpFileObject)
        ObDereferenceObject(lpFileObject);

    if (hFile)
        ZwClose(hFile);

    if (lpBuff)
        ExFreePool(lpBuff);

    KeUnstackDetachProcess(&ApcState);

    if (NT_SUCCESS(ntStatus))
    {
        RtlInitUnicodeString(&usProcessPath, lpFileNameInfo->FileName);

        if (NT_SUCCESS(ntDeviceStatus))
        {
            RtlCopyUnicodeString(lpDosFullName, &usDosDeviceName);
            RtlUnicodeStringCat(lpDosFullName, &usProcessPath);
        }
        else
        {
            RtlCopyUnicodeString(lpDosFullName, &usProcessPath);
        }
    }

    return ntStatus;
}

PEPROCESS FindProcessByName(IN PCWCH lpName)
{
    PEPROCESS lpEprocess = NULL;
    NTSTATUS ntStatus;
    DECLARE_UNICODE_STRING_SIZE(usProcessName, MAX_PATH * sizeof(WCHAR));
    PWCH lpWch;

    for (int i = 8; i < 0x1000000; i += 4)
    {
        ntStatus = Ps2GetProcessFullDosNameByPid((HANDLE)i, &usProcessName);
        if (!NT_SUCCESS(ntStatus)) continue;

        lpWch = usProcessName.Buffer + usProcessName.Length;
        while (lpWch != usProcessName.Buffer && *lpWch != L'\\') lpWch--;
        if (lpWch != usProcessName.Buffer)
        {
            lpWch++;
            if (!_wcsicmp(lpName, lpWch))
            {
                ntStatus = PsLookupProcessByProcessId((HANDLE)i, &lpEprocess);
                break;
            }
        }
    }

    return lpEprocess;
}

#define PROCESS_VM_READ     0x0010
#define PROCESS_VM_WRITE    0x0020

BOOLEAN NTAPI EnumHandleTableRoutine(
    IN PHANDLE_TABLE_ENTRY lpHandleTableEntry,
    IN HANDLE Handle,
    IN PVOID lpEnumParameter
)
{
    UNREFERENCED_PARAMETER(Handle);

    BOOLEAN bRet = FALSE;
    ULONG_PTR lpAddress;
    POBJECT_HEADER lpObjectHeader = NULL;
    ULONG_PTR lpObject;

    do
    {
        if (!MmIsAddressValid(lpHandleTableEntry)) break;

        if (g_dwMajorVersion == 6)
        {
            lpObjectHeader = (POBJECT_HEADER)(((ULONG_PTR)(lpHandleTableEntry->Object)) & 0xfffffffffffffff0);
        }
        else if (g_dwMajorVersion == 10)
        {
            lpAddress = *((PULONG_PTR)lpHandleTableEntry);
            lpObjectHeader = (POBJECT_HEADER)(((LONG_PTR)lpAddress >> 0x10) & 0xfffffffffffffff0);
        }
        else
        {
            break;
        }

        lpObject = (ULONG_PTR)((ULONG_PTR)lpObjectHeader + 0x30);
        if (ObGetObjectType((PVOID)lpObject) == *PsProcessType)
        {
            if (lpObject == (ULONG_PTR)lpEnumParameter)
            {
                lpHandleTableEntry->GrantedAccess &= ~(PROCESS_VM_READ | PROCESS_VM_WRITE);
                bRet = TRUE;
                DBGPRINT("\r\nSuccess\r\n");
                break;
            }
        }

    } while (FALSE);

    return bRet;
}

BOOLEAN Ps2DecreaseProcessHandlePrivilege(IN ULONG dwProcessId, IN PCWCH lpProtectName)
{
    BOOLEAN bRet = FALSE;
    PEPROCESS lpEprocess = NULL;
    PEPROCESS lpProtectedEprocess = NULL;
    NTSTATUS ntStatus;
    ULONG_PTR lpHandleTable = 0;

    UNREFERENCED_PARAMETER(lpProtectName);

    do
    {
        lpProtectedEprocess = FindProcessByName(lpProtectName);
        if (!lpProtectedEprocess)
        {
            ERRPRINT(FindProcessByName, 0);
            break;
        }

        ntStatus = PsLookupProcessByProcessId((HANDLE)dwProcessId, &lpEprocess);
        if (!NT_SUCCESS(ntStatus))
        {
            ERRPRINT(PsLookupProcessByProcessId, ntStatus);
            break;
        }

        ntStatus = PsGetProcessExitStatus(lpEprocess);
        if (ntStatus != STATUS_PENDING)
        {
            DBGPRINT("Process Status != STATUS_PENDING\r\n");
            break;
        }

        if (g_dwMajorVersion == 6)
            lpHandleTable = *((PULONG_PTR)((ULONG_PTR)lpEprocess + 0x200));
        else if (g_dwMajorVersion == 10)
            lpHandleTable = *((PULONG_PTR)((ULONG_PTR)lpEprocess + 0x418));
        else
            break;

        if (lpHandleTable)
        {
            bRet = ExEnumHandleTable((PVOID)lpHandleTable, EnumHandleTableRoutine, lpProtectedEprocess, NULL);
            if (!bRet) break;
        }

        bRet = TRUE;

    } while (FALSE);

    if (lpEprocess)
        ObDereferenceObject(lpEprocess);

    if (lpProtectedEprocess)
        ObDereferenceObject(lpProtectedEprocess);

    return bRet;
}

HANDLE Ob2RegisterCallbacks(
    IN PCWCH lpAltitude,
    IN POBJECT_TYPE* ppObjectType,
    IN POB_PRE_OPERATION_CALLBACK lpPreOperation,
    IN POB_POST_OPERATION_CALLBACK lpPostOperation,
    IN PVOID lpRegistrationContext
)
{
    HANDLE hRegister = NULL;
    OB_CALLBACK_REGISTRATION ObCallbackRegistration = { 0 };
    OB_OPERATION_REGISTRATION ObOperationRegistration = { 0 };

    ObCallbackRegistration.Version = ObGetFilterVersion();
    ObCallbackRegistration.OperationRegistrationCount = 1;
    ObCallbackRegistration.RegistrationContext = lpRegistrationContext;
    RtlInitUnicodeString(&ObCallbackRegistration.Altitude, lpAltitude);
    ObCallbackRegistration.OperationRegistration = &ObOperationRegistration;

    ObOperationRegistration.ObjectType = ppObjectType;
    ObOperationRegistration.Operations = OB_OPERATION_HANDLE_CREATE | OB_OPERATION_HANDLE_DUPLICATE;
    ObOperationRegistration.PreOperation = lpPreOperation;
    ObOperationRegistration.PostOperation = lpPostOperation;

    ObRegisterCallbacks(&ObCallbackRegistration, &hRegister);

    return hRegister;
}

HANDLE Ob2RegisterCallbacks2(
    IN PCWCH lpAltitude,
    IN PVOID lpRegistrationContext,
    IN POB_OPERATION_REGISTRATION lpObOperationRegistration,
    IN USHORT OperationRegistrationCount
)
{
    HANDLE hRegister = NULL;
    OB_CALLBACK_REGISTRATION ObCallbackRegistration = { 0 };

    ObCallbackRegistration.Version = ObGetFilterVersion();
    ObCallbackRegistration.OperationRegistrationCount = OperationRegistrationCount;
    ObCallbackRegistration.RegistrationContext = lpRegistrationContext;
    RtlInitUnicodeString(&ObCallbackRegistration.Altitude, lpAltitude);
    ObCallbackRegistration.OperationRegistration = lpObOperationRegistration;

    ObRegisterCallbacks(&ObCallbackRegistration, &hRegister);

    return hRegister;
}

void Ob2UnRegisterCallbacks(IN HANDLE hRegister)
{
    if (hRegister)
        ObUnRegisterCallbacks(hRegister);
}

OB_PREOP_CALLBACK_STATUS ObPreOperationProcessCallback(
    IN PVOID lpRegistrationContext,
    IN POB_PRE_OPERATION_INFORMATION lpObOperationInformation
)
{
    UNREFERENCED_PARAMETER(lpRegistrationContext);

    PEPROCESS lpEprocess;
    NTSTATUS ntStatus;
    DECLARE_UNICODE_STRING_SIZE(usProcessName, MAX_PATH * sizeof(WCHAR));
    PWCH lpWch;

    do
    {
        if (*PsProcessType != lpObOperationInformation->ObjectType) break;

        lpEprocess = (PEPROCESS)lpObOperationInformation->Object;
        if (!lpEprocess) break;

        ntStatus = Ps2GetProcessFullDosName(lpEprocess, &usProcessName);
        if (!NT_SUCCESS(ntStatus)) break;

        lpWch = usProcessName.Buffer + usProcessName.Length;
        while (lpWch != usProcessName.Buffer && *lpWch != L'\\') lpWch--;
        if (lpWch == usProcessName.Buffer) break;
        lpWch++;
        if (0 != _wcsicmp(L"notepad++.exe", lpWch)) break;

        if (OB_OPERATION_HANDLE_CREATE == lpObOperationInformation->Operation)
        {
            if (lpObOperationInformation->Parameters->CreateHandleInformation.OriginalDesiredAccess & 1)
            {
                lpObOperationInformation->Parameters->CreateHandleInformation.DesiredAccess = 0;
            }
        }

        if (OB_OPERATION_HANDLE_DUPLICATE == lpObOperationInformation->Operation)
        {
            if (lpObOperationInformation->Parameters->DuplicateHandleInformation.OriginalDesiredAccess & 1)
            {
                lpObOperationInformation->Parameters->DuplicateHandleInformation.DesiredAccess = 0;
            }
        }

    } while (FALSE);

    return OB_PREOP_SUCCESS;
}

OB_PREOP_CALLBACK_STATUS ObPreOperationThreadCallback(
    IN PVOID lpRegistrationContext,
    IN POB_PRE_OPERATION_INFORMATION lpObOperationInformation
)
{
    UNREFERENCED_PARAMETER(lpRegistrationContext);

    PEPROCESS lpEprocess;
    NTSTATUS ntStatus;
    DECLARE_UNICODE_STRING_SIZE(usProcessName, MAX_PATH * sizeof(WCHAR));
    PWCH lpWch;

    do
    {
        if (*PsThreadType != lpObOperationInformation->ObjectType) break;

        lpEprocess = (PEPROCESS)IoThreadToProcess((PETHREAD)lpObOperationInformation->Object);
        if (!lpEprocess) break;

        ntStatus = Ps2GetProcessFullDosName(lpEprocess, &usProcessName);
        if (!NT_SUCCESS(ntStatus)) break;

        lpWch = usProcessName.Buffer + usProcessName.Length;
        while (lpWch != usProcessName.Buffer && *lpWch != L'\\') lpWch--;
        if (lpWch == usProcessName.Buffer) break;
        lpWch++;
        if (0 != _wcsicmp(L"notepad++.exe", lpWch)) break;

        if (OB_OPERATION_HANDLE_CREATE == lpObOperationInformation->Operation)
        {
            if (lpObOperationInformation->Parameters->CreateHandleInformation.OriginalDesiredAccess & 1)
            {
                lpObOperationInformation->Parameters->CreateHandleInformation.DesiredAccess = 0;
            }
        }

        if (OB_OPERATION_HANDLE_DUPLICATE == lpObOperationInformation->Operation)
        {
            if (lpObOperationInformation->Parameters->DuplicateHandleInformation.OriginalDesiredAccess & 1)
            {
                lpObOperationInformation->Parameters->DuplicateHandleInformation.DesiredAccess = 0;
            }
        }

    } while (FALSE);

    return OB_PREOP_SUCCESS;
}
